import logging
import warnings

from tqdm import tqdm

from networks.kd_model import NetModel
from networks.kd_model_without_gan import NetModelWithOutGan
from utils.train_options import TrainOptions

warnings.filterwarnings("ignore")
from torch.utils import data
from dataset.datasets import JIDataset, ShandongDataset

args = TrainOptions().initialize()
num_classes = args.classes_num
onehot = num_classes == 1

trainloader = data.DataLoader(ShandongDataset("train", onehot, False), args.batch_size, True, num_workers=4,
                              pin_memory=True)
valloader = data.DataLoader(ShandongDataset("test", onehot, False), 1, False, pin_memory=True)

model = NetModel(args)
best_score = 0
for epoch in range(args.start_epoch, args.epoch_nums):
    for step, data in tqdm(enumerate(trainloader, args.last_step + 1), total=len(trainloader)):
        # model.adjust_learning_rate(args.lr_g, model.G_solver, step)
        # model.adjust_learning_rate(args.lr_d, model.D_solver, step)
        model.set_input(data)
        model.optimize_parameters()
        # model.print_info(epoch, step)
    mean_IU, IU_array = model.evalute_model(model.student, valloader, '0', '512, 512', num_classes, False)
    if IU_array[1] > best_score:
        model.save_ckpt("shandong")
    best_score = max(best_score, IU_array[1])
    logging.info('Epoch: {}, IoU_1:{}, IoU_Best:{}'.format(epoch, IU_array[1], best_score))
